from itertools import product
import os
import json
from create_sub_file.common_file import *
from create_sub_file.write_pbs.write_pbs_common import *
from create_sub_file.write_pbs.write_pbs_amg import *
from create_sub_file.write_slurm.write_slurm_common import *
from create_sub_file.write_slurm.write_slurm_amg import *
from create_sub_file.qsub_pbs.qsub_file import *
from submission.create_sub_file.sbatch_slurm.sbatch_file import *
from create_sub_file.write_slurm.write_slurm_common import *

#TODO: 下面代码还需要优化

def create_amg_loopdict(caselist,run_parameter, processinfo, threasinfo):
    amg_loopdict ={}
    amg_loopdict["numprocess"] = create_process_list(processinfo)
    amg_loopdict["case"] = list(range(len(caselist)))
    amg_loopdict["threads"] = threasinfo
    for para in run_parameter:
        if(isinstance(run_parameter[para],list)):
            amg_loopdict[para] = run_parameter[para]
    # print(amg_loopdict)
    return amg_loopdict

def create_amg_loopindex(amg_loopdict):
    amg_loop_list = list(amg_loopdict.items())
    loop_index = list(range(len(amg_loopdict)))
    for i in range(len(amg_loop_list)):
        loop_index[i] = list(range(len(amg_loop_list[i][1])))
    return amg_loop_list,loop_index

def get_argument_name(tmp_amg_dict,other_runpara):
    except_list = get_except_list()
    argumentname = ""
    for i in tmp_amg_dict:
        if i not in except_list:
            argumentname = argumentname + i[:1] + str(tmp_amg_dict[i]) 
    
    tmp_other = other_runpara.split(" ")
    for i in range(0,len(tmp_other),2):
        tmp_other[i] = tmp_other[i].strip("-")
        argumentname = argumentname + tmp_other[i][:1] + str(tmp_other[i+1])

    return argumentname

def create_amg_subfile_dir(casepath,tmp_amg_dict,other_runpara,program):
    basename = os.path.abspath("./")
    timestr  = get_tiemstr()
    casename = os.path.basename(casepath).split(".")[0]

    nodes  = tmp_amg_dict["nodes"] 
    numprocess = tmp_amg_dict["numprocess"]
    threads = tmp_amg_dict['threads']
    if isinstance(nodes,list):
        nodes = len(nodes)
    processname = "N{0}P{1}T{2}".format(nodes, numprocess, threads)
    
    argumentname = get_argument_name(tmp_amg_dict, other_runpara)

    sub_dir  = os.path.join(basename, casename, processname, argumentname, timestr)
    os.system(f"mkdir -p {sub_dir}")
    for i in program:
        program_name = os.path.basename(i)
        os.system("mkdir -p {0}".format(os.path.join(sub_dir,program_name)))
    return sub_dir

def wtrite_amginfo2json(dir,casename,rhsname,job_scheduler,tmp_amg_dict,other_para, queue, program):
    tmp_dict = {}
    tmp_dict["date"] = get_tiemstr()
    tmp_dict["dir"] = dir
    tmp_dict["casedir"] = casename
    tmp_dict["casename"] = os.path.basename(casename)
    tmp_dict["rhsdir"] = rhsname
    tmp_dict["rhsname"] = os.path.basename(rhsname)
    tmp_dict["subfile"] =  os.path.join(dir, "mysub.{0}".format(job_scheduler))
    tmp_dict["resultlog"] = os.path.join(dir, "sub.result")
    tmp_dict["errolog"] = os.path.join(dir, "sub.erro")
    tmp_dict["nodelist"] =  os.path.join(dir, "nodelist.{0}".format(job_scheduler))
    tmp_dict["queue"] = queue
    tmp_dict["binary"] = program
    tmp_dict["argument"] = {}
    for i in tmp_amg_dict:
        tmp_dict["argument"][i] = tmp_amg_dict[i]
    other_para_list = other_para.split(" ")
    for i in range(0,len(other_para_list),2):
        tmp_dict["argument"][other_para_list[i].strip("-")] = other_para_list[i+1]
    filename = os.path.join(dir, "{0}info.json".format(job_scheduler))
    with open(filename, "w", encoding="utf-8") as fp:
        json.dump(tmp_dict, fp, indent=2, sort_keys=False)

def write_amg_pbs(dir, tmp_amg_dict, myconfig):
    job_scheduler = myconfig["subinfo"]["job_scheduler"]
    filename = os.path.join(dir, "mysub.{0}".format(job_scheduler))
    subinfo = myconfig["subinfo"]
    caselist = myconfig["amg_parameter"]["problem"]["jxpamg_matrix_dir"]
    casename = caselist[tmp_amg_dict['case']]
    casename = os.path.basename(casename).split(".")[0]
    fo = open(filename, "a+" )
    write_pbs_head(fo, subinfo, dir, tmp_amg_dict["nodes"], casename)
    write_pbs_runenvironments(fo, tmp_amg_dict["threads"], dir, myconfig["runenvironments"])
    write_pbs_amg(fo,myconfig["program"],myconfig["amg_parameter"],tmp_amg_dict,myconfig["runenvironments"],dir)
    fo.close()

def write_amg_slurm(dir, tmp_amg_dict, myconfig):
    job_scheduler = myconfig["subinfo"]["job_scheduler"]
    filename = os.path.join(dir, "mysub.{0}".format(job_scheduler))
    subinfo = myconfig["subinfo"]
    caselist = myconfig["amg_parameter"]["problem"]["jxpamg_matrix_dir"]
    casename = caselist[tmp_amg_dict['case']]
    casename = os.path.basename(casename).split(".")[0]
    fo = open(filename, "a+" )
    write_slurm_head(fo, subinfo, dir, tmp_amg_dict["nodes"], casename)
    write_slurm_runenvironments(fo, tmp_amg_dict["threads"], dir, myconfig["runenvironments"])
    write_slurm_amg(fo,myconfig["program"],myconfig["amg_parameter"],tmp_amg_dict,myconfig["runenvironments"],dir)
    fo.close()

def create_sub_amg_file(myconfig):
    amg_parameter_dict = myconfig["amg_parameter"]
    subinfo = myconfig["subinfo"]
    processinfo = subinfo["processinfo"]
    threadsinfo = subinfo["threads"]
    caselist = amg_parameter_dict["problem"]["jxpamg_matrix_dir"]
    rhslist  = amg_parameter_dict["problem"]["rhs_dir"]
    job_scheduler = subinfo["job_scheduler"]

    amg_loopdict = create_amg_loopdict(caselist, amg_parameter_dict["run_parameter"], \
                                        processinfo, threadsinfo)
    amg_loop_list,loop_index = create_amg_loopindex(amg_loopdict)
    
    # 遍历每一种参数组合

    if rhslist == None:
        rhslist = len(caselist)*["None"]

    for para_index in product(*loop_index):
        tmp_amg_dict = {}
        for i in range(len(amg_loop_list)):
            tmp_index = amg_loop_list[i]
            tmp_amg_dict[tmp_index[0]] = \
                tmp_index[1][para_index[i]]
        tmp_amg_dict["nodes"] = get_nodes_info(tmp_amg_dict["numprocess"], subinfo)
        # print(tmp_amg_dict)
        sub_dir = create_amg_subfile_dir(caselist[tmp_amg_dict['case']],tmp_amg_dict,amg_parameter_dict["run_parameter"]["other_runpara"],myconfig["program"])
        wtrite_amginfo2json(sub_dir,caselist[tmp_amg_dict['case']],rhslist[tmp_amg_dict['case']],job_scheduler,\
                            tmp_amg_dict,amg_parameter_dict["run_parameter"]["other_runpara"], subinfo["queue"], myconfig["program"])
        if(job_scheduler == "pbs"):
            write_amg_pbs(sub_dir,tmp_amg_dict,myconfig)
            qsubfile(sub_dir,subinfo)
        elif(job_scheduler == "slurm"):
            write_amg_slurm(sub_dir,tmp_amg_dict,myconfig)
            sbatchfile(sub_dir,subinfo)


